# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import json
import logging
import collections
import pandas as pd
import math
import numpy as np
import datetime

from absl import app, flags

from coin.exchange.base.order_gateway import OrderSide, Order
from coin.experimental.xguo.fruit.base.strategy_sim_base import (
    BaseSimStrategy,
    BaseStrategyConfig,
    convert_book_to_dict,
    convert_trade_to_dict,
)
from coin.experimental.xguo.fruit.util.util import LatestFeed, TimeWindowData
#from coin.experimental.xguo.fruit.cherry.cherry import CherryStrategyConfig

FLAGS = flags.FLAGS


def create_order(order_side, price, qty):
  order = Order(order_id=None,
                product=None,
                price=price,
                qty=qty,
                order_side=order_side,
                order_type=None,
                tag=None,
                accepted=None,
                cancel_sent=None,
                internal=None)
  return order


class BookStats(object):
  def __init__(self, window=30):
    self._data = TimeWindowData(window)
    self._series = pd.Series()
    self._logger = logging.getLogger(__name__)
    self._queue = collections.deque(20)

  @property
  def series(self):
    return self._series

  def update(self, book):
    timestamp = book['timestamp']
    self._data.update(timestamp, book, self._calc)

  def _calc(self, remove_list, new_entry):
    if remove_list:
      time_span = new_entry.value['timestamp'] - remove_list[0].value['timestamp']
      time_span = time_span / 1e9
      if time_span > 35:
        self._logger.error('time span %s', time_span)
        return

      end_price = new_entry.value['data']['mid']
      start_price = remove_list[0].value['data']['mid']
      profit_ratio = (end_price - start_price) / start_price
      self._series[pd.to_datetime(new_entry.timestamp)] = profit_ratio


class TradeStats(object):
  def __init__(self, window=30):
    self._data = TimeWindowData(window)
    self._series = pd.Series()
    self._logger = logging.getLogger(__name__)

  @property
  def series(self):
    return self._series

  def update(self, trade):
    timestamp = trade['timestamp']
    self._data.update(timestamp, trade, self._calc)

  def _calc(self, remove_list, new_entry):
    timestamp = new_entry.value['timestamp']
    if self._series.empty:
      sum_value = 0
    else:
      sum_value = self._series[-1]

    for entry in remove_list:
      sum_value -= entry.value['data']['qty']
    sum_value += new_entry.value['data']['qty']
    self._series[pd.to_datetime(timestamp)] = sum_value


class BookStats1(object):
  def __init__(self, window=30):
    self._data = []
    self._logger = logging.getLogger(__name__)
    self._current_interval = None
    self._window = window
    self._queue = collections.deque(maxlen=5)
    self._latest = 0.

  def update(self, book):
    timestamp = book['timestamp']
    interval_idx = int(math.ceil(timestamp / float(self._window * 1e9)))

    if not self._data:
      self._current_interval = interval_idx
      self._data.append(book['data']['mid'])
    elif self._current_interval == interval_idx:
      end_price = book['data']['mid']
      self._data.append(end_price)
      start_price = self._data[0]
      profit_ratio = (end_price - start_price) / start_price
      self._latest = profit_ratio
    else:
      start_price = self._data[0]
      end_price = self._data[-1]
      profit_ratio = (end_price - start_price) / start_price
      self._queue.append(profit_ratio)
      self._latest = 0.
      self._current_interval = interval_idx
      self._data.clear()

  def __len__(self):
    return len(self._queue)

  def to_array(self):
    return np.array(list(self._queue) + [self._latest])


class TradeStats1(object):
  def __init__(self, window=30):
    self._data = []
    self._logger = logging.getLogger(__name__)
    self._current_interval = None
    self._window = window
    self._queue = collections.deque(maxlen=5)
    self._latest = 0.

  def update(self, trade):
    timestamp = trade['timestamp']
    interval_idx = int(math.ceil(timestamp / float(self._window * 1e9)))

    qty = trade['data']['qty']
    side = trade['data']['side']
    if 'SELL' in side:
      qty = -qty

    if not self._data:
      self._current_interval = interval_idx
      self._data.append(qty)
    elif self._current_interval == interval_idx:
      self._latest = sum(self._data)
    else:
      self._queue.append(sum(self._data))
      self._data.clear()
      self._data.append(qty)
      self._latest = qty

  def __len__(self):
    return len(self._queue)

  def to_array(self):
    return np.array(list(self._queue) + [self._latest])


class CherrySimStrategy(BaseSimStrategy):
  def __init__(self, config):
    super().__init__(config)
    self._latest_feed = LatestFeed()
    self._book_stats = BookStats1(30)
    self._trade_stats = TradeStats1(30)
    self._order = None
    self._ts = None
    self._logger = logging.getLogger(__name__)

  @property
  def product(self):
    return self.product_list[0]

  def on_book(self, book):
    self._ts = book.timestamp
    book_dict = convert_book_to_dict(book, level=10, notional_qty=None)
    self._latest_feed.update_book(book_dict)
    self._book_stats.update(book_dict)
    self._on_feed()

  def on_trade(self, product, trade_event):
    trade_dict = convert_trade_to_dict(product, trade_event)
    self._ts = trade_dict['timestamp']
    self._trade_stats.update(trade_dict)
    self._on_feed()

  def _on_feed(self):
    dt = datetime.datetime.fromtimestamp(self._ts / 1e9)
    x = self._book_stats.to_array()
    x = x / np.linalg.norm(x)
    x = list(x)
    x.reverse()

    y = self._trade_stats.to_array()
    y = y / np.linalg.norm(y)
    c = np.max(np.convolve(x, y))
    data = {
        'ts': self._ts,
        'x': list(x),
        'y': list(y),
        'c': c,
        'tsh': str(dt),
    }
    print(json.dumps(data))


def main(_):
  config = BaseStrategyConfig.from_file(FLAGS.config)
  strategy = CherrySimStrategy(config)
  strategy.run()


if __name__ == '__main__':
  flags.DEFINE_string('config', None, 'Strategy config file.')

  logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  app.run(main)
